import os
import os.path as osp

import cv2
import numpy as np
from tqdm import tqdm


# 棋盘格尺寸：(宽6，高9)
PatternSize = (6, 9)
# 黑方格的大小：17mm
squareSize = 17

def findCorners(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (3, 3), 0)
    ret, corners = cv2.findChessboardCorners(gray, PatternSize, None)
    if ret:
        # 找到角点：提高角点的精确度
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        corners = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
    # 将所有的角点在图像上显示出来
    cv2.drawChessboardCorners(img, PatternSize, corners, ret)
    
    # show
    # cv2.imshow("img", img)
    # cv2.waitKey(0)
    # cv2.imwrite("boardCorners.jpg", img)
    
    if not ret:
        print("can't find corners")

    return ret, corners


def calcObjectPoints():
    # 计算uv空间中角点对应的相机坐标系坐标值，设Z为0
    obj_points = []
    for i in range(PatternSize[1]):
        for j in range(PatternSize[0]):
            obj_points.append([j, i, 0])
    return obj_points


if __name__ == "__main__":
    img_dir = "chess"
    img_list = os.listdir(img_dir)

    objectPoints = []
    cornerPoints = []
    for img_name in tqdm(img_list):
        img_path = osp.join(img_dir, img_name)
        img = cv2.imread(img_path)
        imgSize = img.shape[:2]
        ret, corners = findCorners(img)
        if ret:
            obj_points = np.array(calcObjectPoints(), dtype=np.float32).reshape(-1, 3)
            objectPoints.append(obj_points)
            cornerPoints.append(corners)


    # 执行相机标定
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objectPoints, cornerPoints, imgSize, None, None)
    print("重投影误差(mm):", ret)
    print("内参矩阵:\n", mtx)
    print("畸变系数:\n", dist)
    # print("旋转向量:\n", rvecs)
    # print("平移向量:\n", tvecs)
    with open("cameraMatrix.txt", 'w') as f:
        f.write(str(np.array(mtx).reshape(-1).tolist()) + "\n")
        f.write(str(np.array(dist).reshape(-1).tolist()))
    f.close()


    pass